package com.laotang.第_7_章_堆.力扣_373_查找和最小的;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

public class KSortedArrayMerger {
    /**
     * 多路归并排序
     * 给定 k 个升序的数组，将它们合并成一个有序的数组
     * 例如：
     *  输入： k = 4
     *      a = [1, 2, 6, 9, 10]
     *      b = [2, 3, 7, 11]
     *      c = [4, 8, 9, 13]
     *      d = [5]
     *  输出：
     *      [1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13]
     *
     *  tips：使用小顶堆实现排序
     *
     *  假设每个数组的平均长度为 n
     *  时间复杂度：O(n*k*log(k))
     *  空间复杂度：O(n*k)
     */
    public List<Integer> mergeKSortedArray(List<int[]> data) {
        int k = data.size();
        List<Integer> res = new ArrayList<>();

        // 1. 初始化一个大小为 k 的小顶堆
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(k, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                // 先按照元素值排序，元素值小的放在堆顶
                // 元素值相等的话，按照 元素 所属的数组在 data 中的索引排序，如果小的放在堆顶，这个其实排不排序都无所谓
                int e1 = data.get(o1[0])[o1[1]];
                int e2 = data.get(o2[0])[o2[1]];
                return e1 != e2 ? e1 - e2 : o1[0] - o2[0];
            }
        });

        // 将 k 个数组的第一个元素值都放入到堆中
        for (int i = 0; i < k; i++) {
            if (data.get(i).length == 0) continue;
            minHeap.add(new int[]{i, 0});
        }

        //  合并排序
        while (!minHeap.isEmpty()) {
            // 拿到堆顶最小的元素
            int[] rec = minHeap.poll();
            int i = rec[0], j = rec[1];
            int ele = data.get(i)[j];

            // 将最小元素放入到结果中
            res.add(ele);

            // 处理这个 元素 所在的数组的下一个元素
            j++;
            // 如果没有下一个元素了，则不做任何处理
            if (j == data.get(i).length) continue;
            // 将下一个元素的位置信息放入到堆中
            minHeap.add(new int[]{i, j});
        }

        return res;
    }

    public static void main(String[] args) {
        int[] a = {1, 2, 6, 9, 10};
        int[] b = {2, 3, 7, 11};
        int[] c = {4, 8, 9, 13};
        int[] d = {5};

        List<int[]> data = new ArrayList<>();
        data.add(a);
        data.add(b);
        data.add(c);
        data.add(d);

        List<Integer> res = new KSortedArrayMerger().mergeKSortedArray(data);
        System.out.println(res);
    }

}
